%This script implements test under the assumption that researcher has a different model calibration to produce Figure 6.
%Corresponding author: Rodrigo Adao
%Date: 09/11/2024
%Input: model_data_fgkk.mat, est_shares, output_estimation_step.mat
%Output: Figure 6

%% Preliminaries
clear;
close all;

%set local paths
function_path = "..\Functions\";
addpath(function_path,'-frozen');
set_local_paths;

%Parameter grid
omega_star_grid = [linspace(0,5,11)   , 0*ones(1,11)            ,   0*ones(1,11)            ];
sigma_star_grid = [1.04*ones(1,11)    , 1.04, linspace(1.5,6,10),   1.04*ones(1,11)         ];
sigma_grid      = [2.53*ones(1,11)    , 2.53*ones(1,11)         ,   1.1, 2.53, linspace(4,16,9)];

%Define vector of large varieties in our sample
trade_target = 0.9;

%specify cluster and sample
test_cluster_grid = 1; %no clustering

Nspec = length(sigma_grid);

results_dW = zeros(13,3,Nspec);
results_dM = zeros(13,3,Nspec);
results_dT = zeros(13,3,Nspec);
results_dX = zeros(13,3,Nspec);

for j=1:Nspec

tic
    
%% Import and adjust data

display('------------- running new j -----------------')
    j

%% Import data

display('Implement preliminary steps')

%estimation parameters
test_cluster = test_cluster_grid;

%Model parameters
omega_star = omega_star_grid(j);
sigma_star = sigma_star_grid(j);
sigma      = sigma_grid(j);

%Run preliminary steps for estimation
estimation_preliminary_steps;

clear    rho_* share_pm share_px share_rm delta_ig_0 a_star_ig_0 z_star_ig_0 a_ig_0 aM_g_0 AM_s_0 AD_s_0 betaT_0 beta_s_0 alphaL_s_0 alphaI_s_0 alpha_ks_0 barL_rs_0 Z_rs_0 ...
             m_ig_0 x_ig_0 tau_ig_0 tau_star_ig_0 Lsr Dsg m_val_sample m_val_adj adj_m_s adj_m_gs x_val_sample x_val_adj adj_x_s adj_x_gs

%% IV for welfare

 %Naive IV matrices
    var_adj=sum(ww_dW.^2)/Nobs;
    
    %No GE
    share_NC_nGE_dW  = shareIV_tau;    
    share_wNC_nGE_dW  = ww_dWmat_naive*share_NC_nGE_dW ;
    coef = MC_term1*share_wNC_nGE_dW;
    share_wMC_nGE_dW  = share_wNC_nGE_dW- C*coef;
    share_wMC_nGE_dW  = share_wMC_nGE_dW/var_adj;
    
%Auxiliary matrices for control and welfare adjustments
%load auxiliary adjustment vectors
load(save_share_path + "estimation_shares_t" + trade_target + "_s" + sigma + "_o" + omega_star + "_ss" + sigma_star +  ".mat", 'adjGE_wMC')
    adjGE_wMC_mat= spdiags(adjGE_wMC, 0 , Nobs, Nobs );
    aux = adjGE_wMC_mat*shareMOD_dW;
    share_wMC_dW_dW = aux - C*(MC_term1*aux);
    
    clearvars shareIV_tau aux adjGE_MC_mat coef ww_dWmat ww_dWmat_naive share_NC_nGE_dW share_wNC_nGE_dW adjGE_wMC_mat

%% Implement test

%Predictions 
    dx = shareMOD_dW*shifters;

    coef = MC_term1*dx;
    dx_MC_dW = dx - C*coef;

    coef = MC_term1*dx(sample_naive);
    dxn_MC_dW = dx(sample_naive) - C*coef; 
    
%Outcomes
    dy = dWout_sample;

    coef = MC_term1*dy;
    dy_MC_dW = dy - C*coef;

    coef = MC_term1*dy(sample_naive);
    dyn_MC_dW = dy(sample_naive) - C*coef;   
    
%Difference between outcome and predictions    
    dyt = dy - dx;

    coef = MC_term1*dyt;
    dyt_MC_dW = dyt - C*coef;

    coef = MC_term1*dyt(sample_naive);
    dytn_MC_dW = dyt(sample_naive) - C*coef;   
    
clear C MC_term1 shareMOD_dW coef

%% Welfare stacked specification

disp('running estimation for welfare stacked outcomes')

%welfare
    dW_j = ww_dW'*dx;
    
%Correlation correlation-based tests
    correl_pred = corr([dy_MC_dW, dx_MC_dW]);
    cor_MC_dW = correl_pred(2,1);
    
%Test based on naive IV
    zn_wMC_dW = share_wMC_nGE_dW*shiftersIV;  
    
    [btt_zn_wMC_dW , set_zn_wMC_dW , rjt_zn_wMC_dW , r0t_zn_wMC_dW, pjt_zn_wMC_dW , p0t_zn_wMC_dW ] = implement_test_cluster(dytn_MC_dW, zn_wMC_dW , share_wMC_nGE_dW , shiftersIV, critical, 0, cluster_vec);
    [bty_zn_wMC_dW , sey_zn_wMC_dW , rjy_zn_wMC_dW , r0y_zn_wMC_dW, pjy_zn_wMC_dW , p0y_zn_wMC_dW ] = implement_test_cluster(dyn_MC_dW , zn_wMC_dW , share_wMC_nGE_dW , shiftersIV, critical, 0, cluster_vec);
    [btx_zn_wMC_dW , sex_zn_wMC_dW , rjx_zn_wMC_dW , r0x_zn_wMC_dW, pjx_zn_wMC_dW , p0x_zn_wMC_dW ] = implement_test_cluster(dxn_MC_dW , zn_wMC_dW , share_wMC_nGE_dW , shiftersIV, critical, 0, cluster_vec);

%Test based on welfare GE IV
    zw_wMC_dW = share_wMC_dW_dW*shiftersIV;  
    
    [btt_zw_wMC_dW , set_zw_wMC_dW , rjt_zw_wMC_dW , r0t_zw_wMC_dW, pjt_zw_wMC_dW , p0t_zw_wMC_dW ] = implement_test_cluster(dyt_MC_dW, zw_wMC_dW , share_wMC_dW_dW , shiftersIV, critical, 0, cluster_vec);
    [bty_zw_wMC_dW , sey_zw_wMC_dW , rjy_zw_wMC_dW , r0y_zw_wMC_dW, pjy_zw_wMC_dW , p0y_zw_wMC_dW ] = implement_test_cluster(dy_MC_dW , zw_wMC_dW , share_wMC_dW_dW , shiftersIV, critical, 0, cluster_vec);
    [btx_zw_wMC_dW , sex_zw_wMC_dW , rjx_zw_wMC_dW , r0x_zw_wMC_dW, pjx_zw_wMC_dW , p0x_zw_wMC_dW ] = implement_test_cluster(dx_MC_dW , zw_wMC_dW , share_wMC_dW_dW , shiftersIV, critical, 0, cluster_vec);

%Output
    
    results_zw_dW = full([bty_zw_wMC_dW, btx_zw_wMC_dW, btt_zw_wMC_dW; ...
                          sey_zw_wMC_dW, sex_zw_wMC_dW, set_zw_wMC_dW; ...
                          pjy_zw_wMC_dW, pjx_zw_wMC_dW, pjt_zw_wMC_dW; ...
                          p0y_zw_wMC_dW, p0x_zw_wMC_dW, p0t_zw_wMC_dW]);

    results_zn_dW = full([bty_zn_wMC_dW, btx_zn_wMC_dW, btt_zn_wMC_dW; ...
                          sey_zn_wMC_dW, sex_zn_wMC_dW, set_zn_wMC_dW; ...
                          pjy_zn_wMC_dW, pjx_zn_wMC_dW, pjt_zn_wMC_dW; ...
                          p0y_zn_wMC_dW, p0x_zn_wMC_dW, p0t_zn_wMC_dW]);


    results_cor_dW =full([0             , 0             , cor_MC_dW    ]);

    results_obs_dW = Nobs*ones(1,3);                     
    
    bar = zeros(1,3);
    results_dW_j = [results_zw_dW; bar; results_zn_dW; bar; results_cor_dW; bar; results_obs_dW]; 
    clear results_zw_dW results_zn_dW results_cor_dW results_obs_dW bt* se* rj* r0* pj* p0*
    

%% save output

    results_dW(:,:, j) = results_dW_j;
    results_pred_dW(1,j) = dW_j;

    clear share_wMC_nGE_dW share_wMC_dW_dW zw* zn* *_j

    save(save_estimation_path + "est_Fig_6.mat", 'results_dW', 'results_pred_dW', 'test_cluster_grid', 'omega_star_grid', 'sigma_star_grid', 'sigma_grid', 'trade_target');

toc    
end

%% Organizing results
load(save_estimation_path + "est_Fig_6.mat")

dW_series = full(results_pred_dW' - results_pred_dW(1)');
beta_IV_series = squeeze(results_dW(1,3,:) );
pval_IV_series = squeeze(results_dW(3,3,:) );
rej_95 = pval_IV_series < 0.05;

r1 = rej_95 ;
r0 = rej_95 == 0;

title_size = 14;
label_size = 20;
labely_size = 16;
legend_size = 19;
marker_size = 10;
dot_size =75;
axes_size = 18;

%% Figure 6a
j0 = 1;
jf = 11;

x_par1 = omega_star_grid(j0:jf);
dW_par1 = dW_series(j0:jf);
r1_par1 = r1(j0:jf);
r0_par1 = r0(j0:jf);

dW_par1_r1 =  dW_par1(r1_par1 == 1);
dW_par1_r0 =  dW_par1(r0_par1 == 1);
x_par1_r1  =  x_par1(r1_par1 == 1);
x_par1_r0  =  x_par1(r0_par1 == 1);

figure('DefaultAxesFontSize',axes_size, 'Position', [10 10 600 600]);
plot_omegastar = tiledlayout(1,1);

xmin = 0;
xmax = max(x_par1);

nexttile
scatter(x_par1_r1, dW_par1_r1, 'filled', 'd', 'MarkerFaceColor',[0.5 0.5 0.5], 'SizeData', dot_size)
hold on 
scatter(x_par1_r0, dW_par1_r0, 'filled', 'o', 'MarkerFaceColor','black', 'SizeData', dot_size)
hold on 
xline(0, '-black', 'HandleVisibility','off')
xlim([xmin xmax])
ylim([0 .1])
yticks([0 .02 .04 .06 .08 .1])
legend('rejected', 'not rejected', 'Location', 'southeast', 'FontSize', legend_size)
ylabel('E_t[W (\Delta x (\omega_F) ) - W (\Delta x )]', 'FontSize', label_size)
xlabel('\omega_F', 'FontSize', label_size)

saveas(plot_omegastar, graph_path+ "Fig_6a.png")

saveas(plot_omegastar, graph_path+ "Fig_6a.eps", 'epsc');

%% Figure 6b

j0 = 12;
jf = 22;

x_par1 = sigma_star_grid(j0:jf);
dW_par1 = dW_series(j0:jf);
r1_par1 = r1(j0:jf);
r0_par1 = r0(j0:jf);

dW_par1_r1 =  dW_par1(r1_par1 == 1);
dW_par1_r0 =  dW_par1(r0_par1 == 1);
x_par1_r1  =  x_par1(r1_par1 == 1);
x_par1_r0  =  x_par1(r0_par1 == 1);

figure('DefaultAxesFontSize',axes_size, 'Position', [10 10 600 600]);
plot_sigmastar = tiledlayout(1,1);

xmin = 1;
xmax = max(x_par1);

nexttile
scatter(x_par1_r1, dW_par1_r1, 'filled', 'd', 'MarkerFaceColor',[0.5 0.5 0.5], 'SizeData', dot_size)
hold on 
scatter(x_par1_r0, dW_par1_r0, 'filled', 'o', 'MarkerFaceColor','black', 'SizeData', dot_size)
hold on 
xline(0, '-black', 'HandleVisibility','off')
xlim([xmin xmax])
ylim([-.1 0])
yticks([-.1 -.08 -.06 -.04 -.02 0])
legend('rejected', 'not rejected', 'Location', 'southeast', 'FontSize', legend_size)
ylabel('E_t[W (\Delta x (\sigma_F) ) - W (\Delta x )]', 'FontSize', label_size)
xlabel('\sigma_F', 'FontSize', label_size)

saveas(plot_sigmastar, graph_path + "Fig_6b.png")

saveas(plot_sigmastar, graph_path + "Fig_6b.eps", 'epsc');

%% Plot: sigma

j0 = 23;
jf = 33;

x_par1 = sigma_grid(j0:jf);
dW_par1 = dW_series(j0:jf);
r1_par1 = r1(j0:jf);
r0_par1 = r0(j0:jf);

dW_par1_r1 =  dW_par1(r1_par1 == 1);
dW_par1_r0 =  dW_par1(r0_par1 == 1);
x_par1_r1  =  x_par1(r1_par1 == 1);
x_par1_r0  =  x_par1(r0_par1 == 1);

figure('DefaultAxesFontSize',axes_size, 'Position', [10 10 600 600]);
plot_sigma = tiledlayout(1,1);

xmin = 1;
xmax = max(x_par1);

nexttile
scatter(x_par1_r1, dW_par1_r1, 'filled', 'd', 'MarkerFaceColor',[0.5 0.5 0.5], 'SizeData', dot_size)
hold on 
scatter(x_par1_r0, dW_par1_r0, 'filled', 'o', 'MarkerFaceColor','black', 'SizeData', dot_size)
hold on 
xline(0, '-black', 'HandleVisibility','off')
xlim([xmin xmax])
ylim([-.05 .05])
yticks([-.05 -.03 -.01 .01 .03 .05])
legend('rejected', 'not rejected', 'Location', 'southeast', 'FontSize', legend_size)
ylabel('E_t[W (\Delta x (\sigma) ) - W (\Delta x )]', 'FontSize', label_size)
xlabel('\sigma', 'FontSize', label_size)

saveas(plot_sigma, graph_path + "Fig_6c.png")

saveas(plot_sigma, graph_path + "Fig_6c.eps", 'epsc');
